import tqdm
import onnx
import onnxruntime
import numpy as np
from pathlib import Path


def onnx_predict(onnx_path, input_dir, output_dir):

    output_dir = Path(output_dir)
    output_dir.mkdir(parents=True, exist_ok=True)

    # 检查模型并获取session
    onnx_model = onnx.load(onnx_path)
    onnx.checker.check_model(onnx_model)
    ort_session = onnxruntime.InferenceSession(onnx_path)
    input_shape = ort_session.get_inputs()[0].shape
    input_shape[0] = -1

    for input_bin in tqdm.tqdm(Path(input_dir).iterdir(), desc="Processing"):
        # 准备输入数据
        input_data = np.fromfile(input_bin.__str__(), dtype=np.float32)
        input_data = input_data.reshape(input_shape)
        ort_inputs = {ort_session.get_inputs()[0].name: input_data}

        # 推理并保存结果
        ort_outs = ort_session.run(None, ort_inputs)
        output_bin = output_dir / input_bin.name
        ort_outs[0].tofile(output_bin)

    print("Done.")

if __name__ == "__main__":

    import argparse
    parser = argparse.ArgumentParser(description='T2T-ViT Validate.')
    parser.add_argument('--onnx-path', type=str, metavar='PATH', help='')
    parser.add_argument('--in-dir', type=str, metavar='DIR', help='path to dataset')
    parser.add_argument('--out-dir', type=str, metavar='PATH', help='')
    args = parser.parse_args()

    onnx_predict(
        args.onnx_path,
        args.in_dir,
        args.out_dir,
    )